
from app.libs.enums import (
    AssetTaskGroup,
    TaskTriggerType,
    WS_TASK_GROUP_TUPLE,
    TaskType
)
from app.db.models.asset_tasks import AssetTasks
from app.db.models.users import User
from app.db.models.jobs import (
    Job,
    VulSettings,
    SslSettings,
    SecurityEventSettings,
    ContentSettings,
    ChangeCheckSettings
)
from app.libs.utils import enable_task_group


class UpdateSettings:

    def __init__(
        self,
        task_types=None,
        username=None,
        uid=None,
        job_id=None,
        target_url=None,
        source_ip=None,
    ):
        """
        复制周期web安全监测任务配置：
        1、在同一个用户里， 把一个资产的周期性web监测任务配置参数同步到其他资产的周期性web监测任务
        2、只会更新已经存在的web安全监测任务， 不会新建

        :job_id 如果job_id=None， 则target_url与source_ip必须同时传入。
        特别情况，资产没有源站IP时，source_ip传入'null'
        :task_types 需要同步的web安全监测项, 默认所有监测项，多个以逗号分割， eg: ssl,vul
        """
        self.user = None
        if uid:
            self.user = User.objects.filter(id=uid).first()
        elif username:
            self.user = User.objects.filter(username=username).first()
        assert self.user is not None, "用户信息错误"
        self.base_job = None
        job_query = {"uid": self.user.id}
        assert job_id or target_url, "模版资产信息错误"
        if job_id:
            job_query.update({"id": job_id})
        else:
            job_query.update({"targetUrl": target_url})
            if source_ip == "null":
                job_query.update({"sourceIp__in": ['', None]})
            else:
                job_query.update({"sourceIp": source_ip})
        jobs = Job.objects.filter(**job_query)
        assert jobs.count() == 1, "模版资产信息错误"
        self.base_job = jobs.first()
        self.base_asset_task = AssetTasks.objects.filter(
            jobId=self.base_job.id,
            triggerType=TaskTriggerType.schedule.value,
            taskGroupType=AssetTaskGroup.web_security.value
        ).first()
        assert self.base_asset_task, "模版资产没有周期性web安全监测任务"
        self.task_types = self.parse_task_types(task_types=task_types)

    def parse_task_types(self, task_types: str):
        if task_types is None:
            task_types = list(WS_TASK_GROUP_TUPLE)
        else:
            task_types = task_types.strip().split(",")
        assert isinstance(task_types, list)
        task_types = [t for t in task_types if t in WS_TASK_GROUP_TUPLE]
        # 模版资产不支持ssl监测是， 不复制ssl安全任务配置
        if TaskType.ssl.value in task_types and not self.base_job.targetUrl.startswith("https"):
            task_types.remove(TaskType.ssl.value)
        print(f"更新任务类型: {task_types}")
        assert task_types, "更新任务类型错误"
        return task_types

    def update_targets(self, targets: str):
        """
        更新到固定资产web安全监测任务
        :targets  多个以逗号分割，没有源站IP时，源站IP输入null, eg: "targetUrl1|sourceIp1,targetUrl2|null"
        """
        assert targets and isinstance(targets, str), "目标资产信息错误"
        targets = targets.split(",")
        job_query = {"uid": self.user.id}
        target_query = []
        for target in targets:
            target_url, source_ip = target.split("|")
            tq = {"targetUrl": target_url}
            if source_ip == 'null':
                tq.update({"sourceIp": {"$in": ['', None]}})
            else:
                tq.update({"sourceIp": source_ip})
            target_query.append(tq)
        job_query.update({"$or": target_query})
        job_ids = Job._get_collection().distinct("_id", job_query)
        asset_tasks = AssetTasks.objects.filter(
            uid=self.user.id,
            jobId__in=job_ids,
            triggerType=TaskTriggerType.schedule.value,
            taskGroupType=AssetTaskGroup.web_security.value
        )
        self.update_task_settings(asset_tasks)

    def update_all_targets(self):
        """
        更新到用户所有资产web安全监测任务
        """
        asset_tasks = AssetTasks.objects.filter(
            uid=self.user.id,
            jobId__ne=self.base_job.id,
            triggerType=TaskTriggerType.schedule.value,
            taskGroupType=AssetTaskGroup.web_security.value
        )
        self.update_task_settings(asset_tasks)

    def marge_vul_settings(self, asset_task):
        base_settings = getattr(self.base_asset_task, "vulSettings", None)
        if (not base_settings) or (not base_settings.taskId):
            return asset_task
        base_settings = base_settings.to_dict()
        settings = asset_task.vulSettings.to_dict()
        base_collect = base_settings.get("collect")
        collect = settings.get("collect")
        collect.update({
            "depth": base_collect.get("depth"),
            "maxPage": base_collect.get("maxPage"),
            "type": base_collect.get("type"),
            "vul": base_collect.get("vul"),
            "speed": base_collect.get("speed")
        })

        settings.update({
            "collect": collect,
            "enableSitePortraitTriggerMonitor": base_settings.get("enableSitePortraitTriggerMonitor"),
            "alert": base_settings.get("alert"),
            "vulType": base_settings.get("vulType")
        })
        asset_task.vulSettings = VulSettings.from_dict(settings)
        return asset_task

    def marge_ssl_settings(self, asset_task):
        base_settings = getattr(self.base_asset_task, "sslSettings", None)
        if (not base_settings) or (not base_settings.taskId):
            return asset_task
        base_settings = base_settings.to_dict()
        settings = asset_task.sslSettings.to_dict()
        base_collect = base_settings.get("collect")
        collect = settings.get("collect")
        collect.update({
            "pluginCategories": base_collect.get("pluginCategories")
        })
        settings.update({
            "collect": collect,
            "enableSitePortraitTriggerMonitor": base_settings.get("enableSitePortraitTriggerMonitor"),
            "alert": base_settings.get("alert"),
        })
        asset_task.sslSettings = SslSettings.from_dict(settings)
        return asset_task

    def marge_securityEvent_settings(self, asset_task):
        base_settings = getattr(self.base_asset_task, "securityEventSettings", None)
        if (not base_settings) or (not base_settings.taskId):
            return asset_task
        base_settings = base_settings.to_dict()
        settings = asset_task.securityEventSettings.to_dict()
        base_collect = base_settings.get("collect")
        collect = settings.get("collect")
        collect.update({
            "depth": base_collect.get("depth"),
            "maxPage": base_collect.get("maxPage"),
            "plugins": base_collect.get("plugins"),
            "speed": base_collect.get("speed"),
        })
        settings.update({
            "collect": collect,
            "enableSitePortraitTriggerMonitor": base_settings.get("enableSitePortraitTriggerMonitor")
        })
        asset_task.securityEventSettings = SecurityEventSettings.from_dict(settings)
        return asset_task

    def marge_content_settings(self, asset_task):
        base_settings = getattr(self.base_asset_task, "contentSettings", None)
        if (not base_settings) or (not base_settings.taskId):
            return asset_task
        base_settings = base_settings.to_dict()
        settings = asset_task.contentSettings.to_dict()
        base_collect = base_settings.get("collect")
        collect = settings.get("collect")
        collect.update({
            "customKeywordTypes": base_collect.get("customKeywordTypes"),
            "depth": base_collect.get("depth"),
            "maxPage": base_collect.get("maxPage"),
            "plugins": base_collect.get("plugins"),
            "privacyDisclosureTypes": base_collect.get("privacyDisclosureTypes"),
            "systemKeywordTypes": base_collect.get("systemKeywordTypes"),
        })
        settings.update({
            "collect": collect,
            "enableSitePortraitTriggerMonitor": base_settings.get("enableSitePortraitTriggerMonitor")
        })
        asset_task.contentSettings = ContentSettings.from_dict(settings)
        return asset_task

    def marge_changeCheck_settings(self, asset_task):
        base_settings = getattr(self.base_asset_task, "changeCheckSettings", None)
        if (not base_settings) or (not base_settings.taskId):
            return asset_task
        base_settings = base_settings.to_dict()
        settings = asset_task.changeCheckSettings.to_dict()
        base_collect = base_settings.get("collect")
        collect = settings.get("collect")
        collect.update({
            "baseChangeCount": base_collect.get("baseChangeCount"),
            "changeType": base_collect.get("changeType"),
            "depth": base_collect.get("depth"),
            "enableBaseChange": base_collect.get("enableBaseChange"),
            "maxPage": base_collect.get("maxPage"),
        })
        settings.update({
            "collect": collect,
        })
        asset_task.changeCheckSettings = ChangeCheckSettings.from_dict(settings)
        return asset_task

    def update_task_settings(self, asset_tasks):
        for asset_task in asset_tasks:
            for task_type in self.task_types:
                getattr(self, f"marge_{task_type}_settings")(asset_task)
            asset_task.save()
            if asset_task.enableMonitor:
                enable_task_group(asset_task)
            print(f"更新成功 {asset_task.id}: {asset_task.targetUrl}")


def update_task_settings(
    task_types=None,
    username=None,
    uid=None,
    job_id=None,
    target_url=None,
    source_ip=None,
    is_all=False,
    targets=None
):
    print("开始同步web安全监测任务配置......")
    obj = UpdateSettings(
        task_types=task_types,
        username=username,
        uid=uid,
        job_id=job_id,
        target_url=target_url,
        source_ip=source_ip,
    )
    if is_all:
        obj.update_all_targets()
    elif targets:
        obj.update_targets(targets)
    print("同步web安全监测任务配置完成......")
